#!/usr/bin/python
##################################################################################
# Author: Huaqin Xu (huaxu@ucdavis.edu)
# Supervisor: Alexander Kozik (akozik@atgc.org)
# Date: Sep. 11. 2009
# Description:
#
# This python script count the coverage of blast file. 
#
# =================================================================================
# input arguments:
#	1.blast all-hits file.
#	2.Configure file
#	3.Option: 0 -- draw image and hits table; others -- draw image
#		  
# Output: blast coverage table and image file.
#
######################################################################################

import Image,ImageDraw, ImageFont
import sys
import re
import array
import os
from math import *
from os.path import exists, join, basename, splitext


# ---------------------------functions ------------------------------------------------
# ---------------- Open and read file functions ---------------------------------------
def open_file(file_name, mode):
	if file_name == "":
		print 'Empty input file name!'
		raw_input("\nPress the enter key to exit.")
		sys.exit(0)

	try:
		the_file = open(file_name, mode)
	except(IOError), e:
		print "Unable to open the file", file_name, "Ending program.\n", e
		raw_input("\nPress the enter key to exit.")
		sys.exit(0)
	else:
		return the_file

def read_file(afile):
	try:
		flines = afile.readlines()
	except:
		print 'Failed to read from: ', afile
		sys.exit(0)
	else:
		return flines


# ------------------- get Configurate parameter --------------------------------------
def setInt(t):
	for i in range(len(t)):
		t[i]=int(t[i])
	return t

def getConfig(configlines):
	global configset

	configset['alignment'] = "ALL" # PRM + ALT
	configset['identity'] = "ALL" # percent of identity
	configset['perfectalign_length'] = "ALL" 
	configset['alignment_length'] = "ALL" # length of alignment
	configset['full_length'] = "N" 
	configset['pos_color'] = 255, 0, 255
	configset['neg_color'] = 0, 255, 255
	configset['all_color'] = 0, 0, 255
	configset['print_label'] = 'Y'


	if(len(configlines) == 0):
		print "Empty configuration file!"
		exit(0)

	for l in configlines:

		if l != "\n" and l[0] != "#":
			set=l.strip().split("=")
			avalue=set[1].strip()
			if(avalue != ''):		
				if(avalue.count(',')>0):
					avalue=tuple(setInt(avalue.split(',')))
				configset[set[0].strip()]=avalue
	print configset


# ------------------- get Data to be displyed------------------------------------
def getData(lines):
	global query
	global pos
	global neg
	global delimiter
	global configset
	
	delimiter="\t"
	datalen=len(lines)
	if(datalen == 0):
		print "Empty Data file!"
		sys.exit(0)
	else:
		colcount=lines[0].count(delimiter)+1
	query=[]	
	pos={}
	neg={}
	
	aq = 'x'
	for l in range(0, datalen):

	# check for empty lines and incorrect field numbers
		if lines[l] != '\n':
			if lines[l].count(delimiter)==colcount-1 and colcount>4:
				arow=lines[l].rstrip().split(delimiter)
				cond=1
				alen=int(arow[15])
				if(arow[0] != aq):
					query.append(arow[0])
					aq = arow[0]
					pos[aq]= [0]*alen;
					neg[aq]= [0]*alen;					
				
				if(configset['alignment'] != 'ALL' and configset['alignment'] != arow[8]):
					cond = 0
#					print "1 "+str(cond)
				elif(configset['identity'] != 'ALL' and configset['identity'] != arow[3]):
					cond = 0
#					print "2 "+str(cond)
				elif(configset['perfectalign_length'] != 'ALL' and not (int(arow[4]) >= configset['perfectalign_length'][0]  and int(arow[4]) <= configset['perfectalign_length'][1] )):
					cond = 0
#					print "3 "+str(cond)
				elif(configset['alignment_length'] != 'ALL' and not (int(arow[5]) >= configset['alignment_length'][0]  and int(arow[5]) <= configset['alignment_length'][1] )):
					cond = 0
#					print "4 "+str(cond)
				elif(configset['full_length'] != 'N' and arow[5] != arow[16]):
					cond = 0
#					print "5 "+str(cond)
				
				if(cond !=0):
					start = int(arow[10])-1
					end = int(arow[11])
		
					for m in range(start,end):
						if(arow[9] == '+/+'):
							pos[aq][m] = pos[aq][m]+1
						else:
							neg[aq][m] = neg[aq][m]+1					
				
						
			else:
				print "Error: Line #%s has inconsistent number of columns.\n" %(l+1)
				if(l != datalen-1):
					colcount =lines[l+1].count(delimiter)+1
		else:
			print "Skip an empty line at #%s.\n" %(l+1)



# ------------------- print table -------------------
def printtable():
	global query
	global pos
	global neg
	global delimiter
	
	for q in query:
		alen=len(pos[q])
		if(alen !=0):
			line = "Seq_ID" + delimiter + "forawrd hits" + delimiter + "reverse hits" + delimiter + "Total\n"
			outf.write(line)
			for c in range(alen):
				total = pos[q][c] + neg[q][c]
				line = q + delimiter + str(pos[q][c])+ delimiter + str(neg[q][c]) + delimiter + str(total) + '\n'
				outf.write(line)
		else:
			print "No hits for"+q
			

# --------------------------------- draw image  ------------------------------------
def drawTitle(title, imgformat="PNG", fgcolor=(0,0,0), bgcolor=(255,255,255)):

	font = ImageFont.load_default()
	img = Image.new("RGBA", (20,49), bgcolor)
    
	# calculate space needed to render text
	draw = ImageDraw.Draw(img)
	sizex, sizey = draw.textsize(title, font=font)
	img = img.resize((sizex,sizey))
	draw = ImageDraw.Draw(img)
    
	# render label into image draw area
	draw.text((0, 0), title, fill=fgcolor, font=font)
	imgOut = Image.new("RGBA", (sizex, sizey), bgcolor)
	draw = ImageDraw.Draw(imgOut)

	return img
	

def drawImg(q):
	global pos
	global neg
	global infile
	global configfile
	global configset

	
	imageY=500
	imageX=(len(pos[q]))*2+10
	midX=250
	ftcolor=(0,0,0)
	bgcolor=(255,255,255)
	poscolor=configset['pos_color']
	negcolor=configset['neg_color']
	greycolor=(196,196,196)
	allcolor=configset['all_color']
	
	img = Image.new("RGBA", (imageX,imageY), bgcolor)
	draw = ImageDraw.Draw(img)
	
	startX=5;
	endX=imageX-5;
	
	titleimg1 = drawTitle(q)
	img.paste(titleimg1, (startX+10, 10))
	
	title2 = "[+/+]" + str(sum(pos[q]))
	titleimg2 = drawTitle(title2)
	img.paste(titleimg2, (endX-100, 10))
	title3 = "[+/-]" + str(sum(neg[q]))
	titleimg3 = drawTitle(title3)
	img.paste(titleimg3, (endX-100, 20))
	title4 = "[ALL]" + str(sum(pos[q])+sum(neg[q]))
	titleimg4 = drawTitle(title4)
	img.paste(titleimg4, (endX-100, 30))
	title5 = basename(infile)
	titleimg5 = drawTitle(title5)
	img.paste(titleimg5, (startX+10, imageY-20))
	title6 = basename(configfile)
	titleimg6 = drawTitle(title6)
	img.paste(titleimg6, (endX-150, imageY-20))
	
	draw.line([(startX,imageY/2),(endX,imageY/2)],fill=ftcolor)
	

	per = max(pos[q])+max(neg[q])
	for c in range(len(pos[q])):
		posY = pos[q][c]*200/per
		negY = neg[q][c]*200/per
		
		draw.line([(startX+c*2, midX), (startX+c*2, midX-posY)], fill=poscolor)
#		draw.line([(startX+c*2, midX-posY), (startX+c*2, midX-posY-negY)], fill=greycolor)
		draw.point((startX+c*2, midX-posY-negY), fill=allcolor)
		draw.line([(startX+c*2, midX), (startX+c*2, midX+negY)], fill=negcolor)		
#		draw.line([(startX+c*2, midX+negY), (startX+c*2, midX+negY+posY)], fill=greycolor)
		draw.point((startX+c*2, midX+negY+posY), fill=allcolor)
			
	if(configset['print_label'] == 'Y'):
		for c in range(len(pos[q])):
			if(c % 100 == 0):
				draw.line([(startX+c*2, midX+5), (startX+c*2, midX-5)], fill=ftcolor)
				draw.text((startX+c*2, midX-10), str(c), fill=ftcolor)
			
	return img

#----------------------------- main ------------------------------------------------------

# ----- get options and file names and open files -----
if len(sys.argv) == 4:
	infile=sys.argv[1]
	configfile=sys.argv[2]
	choice=sys.argv[3]
else: 
	print len(sys.argv)
	print 'Usage: [1]all-hits file, [2]configure file [3]choice'
	sys.exit(1) 

configset={}
configf=open_file(configfile,'r')
configflines = read_file(configf)
getConfig(configflines)

inf=open_file(infile,'r')
inlines = read_file(inf)
infilebase = splitext(basename(infile))[0].split('-')[0]
#infilebase = splitext(basename(infile))[0]
configbase = splitext(basename(configfile))[0]


getData(inlines)

if choice == '1':
	for q in query:
		if(max(pos[q]) !=0 or max(neg[q]) !=0):
			imgfile = q + '_FROM_' + infilebase +"_"+ configbase + '.png'
			img=drawImg(q)
			img.save(imgfile, "PNG")
		else:
			print "No image for " +q +" at " +configbase	
elif choice == '2':
	outfile= infilebase +"_"+ configbase+ '.out'
	outf=open_file(outfile,'w')	
	printtable()
else:
	outfile= infilebase +"_"+ configbase+ '.out'
	outf=open_file(outfile,'w')	
	printtable()
	for q in query:
		if(max(pos[q]) !=0 or max(neg[q]) !=0):
			imgfile = q + '_FROM_' + infilebase +"_"+ configbase + '.png'
			img=drawImg(q)
			img.save(imgfile, "PNG")
		else:
			print "No image for " +q +" at " +configbase
		
		